%-------------------------------------------------------------------------------
% $Id: heat2d_finite_diff.m,v 1.5 2011/07/05 06:26:09 paul Exp $
% $Date: 2011/07/05 06:26:09 $
% $Author: paul $
%-------------------------------------------------------------------------------
% Finite Differences Solution to the 2D Heat Equation
% Arguments:  M (int):          number of points on mesh x and y axis
%             N (int):          number of time steps
%             method (string):  'explicit' (default) or 'implicit'
%             pattern (string): 'square', 'circles', 'cross' (initial heat dist)
%             graphic (string): 'surf', 'imagesc', 'mesh'
%             solver (string):  'inversion', 'gauss_seidel' 
%
% Heat Equation Defn:  du{t} = alpha^2[du{xx} + du{yy}]  
% Domain:              -1 =< x,y =< +1, 0 =< t =< 1
% Assumptions:         alpha = 1.0
% Initial Conditions:  u(x,y,0) = 1 if |x|<0.5 and |y|<0.5,
%                      u(x,y,0) = 0 elsewhere
% Boundary Conditions: u(-1,y,t) = u(1,y,t) = u(x,-1,t) = u(x,1,t) = 0
%-------------------------------------------------------------------------------
function heat2d_finite_diff(M, N, method, pattern, graphic, solver)
if nargin < 2                  % use all default parameters 
  M = 25;    
  N = 200;  
  method = 'implicit';
  graphic = 'surf';
  pattern = 'square';
  solver = 'gauss_seidel';
elseif nargin < 3              % use default M and N 
  method = 'implicit';  
  graphic = 'surf';
  pattern = 'square';
  solver = 'gauss_seidel';  
elseif nargin < 4  
  pattern = 'square';
  graphic = 'surf';
  solver = 'gauss_seidel';  
elseif nargin < 5  
  graphic = 'surf';  
  solver = 'gauss_seidel';  
elseif nargin < 6
  solver = 'gauss_seidel';    
end
msg = sprintf('Two Dimensional Heat Equation Finite Differences Solution');
disp(msg);
msg = sprintf('Arguments and Options: M, N, method, pattern, graphic, solver');
disp(msg);
msg = sprintf('           M x M - dimesions of 2D grid');
disp(msg);
msg = sprintf('           N     - number of timesteps');
disp(msg);
msg = sprintf('           method - implicit, explicit');
disp(msg);
msg = sprintf('           pattern - square, cross, circles');
disp(msg);
msg = sprintf('           graphic - surf, mesh, imagesc');
disp(msg);
msg = sprintf('           solver - inversion, gauss_seidel (implicit method)');
disp(msg);
msg = sprintf(' ');
disp(msg);

msg = sprintf('INFO:\tArgument Values:');
disp(msg);
msg = sprintf('INFO:\tM=%d N=%d method=%s pattern=%s',M,N,method,pattern);
disp(msg);    
msg = sprintf('INFO:\tplot=%s solver=%s',graphic,solver);
disp(msg);    

% done with argument processing
%-------------------------------------------------------------------------------

T = 1.0;                       % time interval (0,T) 
alpha = 1.0;                      
length = 2*1.0;                % sides x and y: -1 =< x,y =< +1

h = length / M;                % dx
k = T / N;                     % dt
lamda = (alpha^2) * k/(h^2);   % lamda =< 0.25 for stability in
                               % explicit method

msg = sprintf('INFO:\tParameters for the 2D Finite Difference Solution:');
disp(msg); 
msg = sprintf('INFO:\tT=%f alpha=%f h=%f k=%f lamda=%f',T,alpha,h,k,lamda);
disp(msg); 
prompt = sprintf('INPUT:\tPress return to continue');
response=input(prompt);

if (lamda > 0.25 & strcmp(method,'explicit'))
  warning = sprintf('WARNING:Explicit PDE solution unstable if lamda>0.25');
  disp(warning);
  hint = sprintf('INFO:\tAdjust parameters M, N, and T for convergence');
  disp(hint);
  prompt = sprintf('INPUT:\tPress return to continue');
  response=input(prompt);
end

%-------------------------------------------------------------------------------
% Setup AVI file and parameters
aviobj = avifile('heat2d_finite_diffs.avi');
aviobj.compression = 'None';   % no other option for Unix
aviobj.fps = floor(N/30);      % frames per second (scaled to 30 seconds)
aviobj.quality = 10;           % low quality is OK

% optimize iteration loop by using function handles to preselect plot function
if strcmp(graphic,'imagesc')
  do_plot = str2func('do_imagesc');
else
  if strcmp(graphic,'surf')
      do_plot = str2func('do_surf');
  else
      do_plot = str2func('do_mesh');
  end
end

%-------------------------------------------------------------------------------
% Animation loop
fig = figure;

lamda4 = 4 * lamda;                % precompute values outside of loop!
neg_lamda = -lamda;
factor_explicit = 1 - lamda4;           
factor_implicit = 1 + lamda4;

%solver = 'gauss_seidel';  

mask = make_mask(M);               % mask sets the boundary conditions
z = zeros(M+1);                    % initialize frame and delayed versions
z1 = z;

tic;                               % timing stopwatch on

if strcmp(method,'explicit')  
  for j=1:N
    if (j==1)                      % set up initial conditions of vibration
      x = linspace(-1,1,M+1);        
      y = linspace(-1,1,M+1);            
      z = make_init(M,pattern);
    else                           % remaining time iterations
      z = zeros(M+1);    
      z(2:M+1,1:M+1) = z1(1:M  ,1:M+1);                   % shift i+1
      z(1:M  ,1:M+1) = z(1:M  ,1:M+1) + z1(2:M+1,1:M+1);  % shift i-1    
      z(1:M+1,1:M)   = z(1:M+1,1:M)   + z1(1:M+1,2:M+1);  % shift j+1
      z(1:M+1,2:M+1) = z(1:M+1,2:M+1) + z1(1:M+1,1:M);    % shift j-1    
      z = (factor_explicit*z1 + lamda*z) .* mask;        
    end
    z1 = z;  % update delayed versions of membrane heat state 
    do_plot(x,y,z);    % displaying image with selected plot function
    title('2D Heat Equation Finite Differences Solution');
    
    frame = getframe(fig);  % NOTE: this call takes half of total program
                            % execution time based on profile report
    aviobj = addframe(aviobj,frame);
  end 
else % implicit method
  for j=1:N
    if (j==1)                      % set up initial conditions of vibration    
      x = linspace(-1,1,M+1);        
      y = linspace(-1,1,M+1);  
      z = make_init(M,pattern);
      zv = reshape(z,(M+1)^2,1);   % convert matrix to a 1D vector
      % pentadiagonal matrix describes implicit PDE transfer function
      b = sparse_pentadiagonal(factor_implicit, neg_lamda, neg_lamda, ...
                           neg_lamda, neg_lamda, M+1, (M+1)^2);
    else
      if strcmp(solver,'inversion')      
        zv = b\zv;            % simple matrix inversion        
      else
        zv = gauss_seidel(b,zv);
      end
      z = reshape(zv,M+1,M+1) .* mask; % 2D for BC masking and plotting
      zv = reshape(z,(M+1)^2,1);   % convert back to 1D
    end      
    do_plot(x,y,z);    % displaying image with selected plot function
    title('2D Heat Equation Implicit Finite Differences Solution');
    frame = getframe(fig);
    aviobj = addframe(aviobj,frame);
  end
end

time_elapsed = toc;                % timing stopwatch off
msg = sprintf('INFO:\tElapsed Runtime: %d seconds', floor(time_elapsed));
disp(msg);

%-------------------------------------------------------------------------------
% Clean up rendering frame and AVI file

close(fig);
aviobj = close(aviobj);

%-------------------------------------------------------------------------------
% Utility functions
%-------------------------------------------------------------------------------
function init = make_init(M,pattern)
  init = zeros(M+1);
  M2=floor((M+1)/2);  
  if strcmp(pattern,'circles')
    x = linspace(-1,1,M2); 
    for i=1:M2
      for j=1:M2
        if x(i)^2+x(j)^2 < 1
          init(M2+i,M2+j) = 1.0;
          init(i,j) = 1.0;          
        end
      end
    end
  elseif strcmp(pattern,'cross')    
    init(:,M2-1:M2+1) = 1.0;
    init(M2-1:M2+1,:) = 1.0;    
  else
    % quarter square in the center of the surface    
    width = M2;    
    low = ceil((M+1-width)/2)+1;
    high = low+width;
    init(low:low+width-1,low:low+width-1) = ones(width);  
  end
  
function mask = make_mask(M, type)
  % construct square boundary mask with the edges set to zero
  mask = zeros(M+1);                 
  mask(2:M,2:M) = ones(M-1);         

function do_imagesc(x,y,z)
  imagesc(z);
  
function do_surf(x,y,z)  
   surf(x,y,z); 
   axis([-1 1 -1 1 0 1]);
   daspect([1 1 1]);
   view(30,30); % view at 30 deg azimuth, 30 deg elevation
   colormap hot;
   colorbar;
   grid off;
   % Fancy surface plot setup
   camlight left;    % light source for Phong shading
   lighting phong;

function do_mesh(x,y,z)  
   mesh(x,y,z); 
   axis([-1 1 -1 1 0 1]);
   daspect([1 1 1]);
   view(30,30); % view at 30 deg azimuth, 30 deg elevation

function m = sparse_pentadiagonal(m,p1,m1,p2,m2,offset,n)
% construct sparse pentadiagonal matrix
a = zeros(n);
% columns of b represent diagonal bands in constructed matrix
b = [repmat(m2,n,1) repmat(m1,n,1) repmat(m,n,1) repmat(p1,n,1) repmat(p2,n,1)];
d = [-offset, -1, 0, 1, offset]; % diagonal band location specifier
m = spdiags(b,d,a);

function m = full_pentadiagonal(m,p1,m1,p2,m2,offset,n)
% Construct a full pentadiagonal matrix by creating 
% template diagonals using shifted eye matrices
% These templates are scaled and then added together
% NOTE - this method is only practical for small (500x500) 
% matrices. Large problems need sparse matrices for performance!!

diag_main = eye(n);
diag_p1 = [zeros(n,1) diag_main(:,1:n-1)];
diag_m1 = [diag_main(:,2:n) zeros(n,1)];
diag_p2 = [zeros(n,offset) diag_main(:,1:n-offset)];
diag_m2 = [diag_main(:,offset:n) zeros(n,offset-1)];

diag_main_scaled = m .* diag_main;
diag_p1_scaled = p1 .* diag_p1;
diag_m1_scaled = m1 .* diag_m1;
diag_p2_scaled = p2 .* diag_p2;
diag_m2_scaled = m2 .* diag_m2;

m = diag_main_scaled + diag_p1_scaled + diag_m1_scaled;
m = m + diag_p2_scaled + diag_m2_scaled;

function x = gauss_seidel(a,b)
% Gauss-Seidel method for solving b=ax
% x=a\b is pretty good in Matlab now that banded matrices are
% automatically recognized and handled. Try this method anyway.
% a is an MxM matrix
% x,b are Mx1 vectors

[m n] = size(a);

l = tril(a,-1);    % lower triangular matrix 
u = triu(a,1);     % upper triangular matrix
d = diag(diag(a)); % diagonal matrix

% inverses of triangular matrices are trivial
ld_inv = inv(l+d);
ld_inv_b = ld_inv * b;
ld_inv_u = ld_inv * u;

x= zeros(m,1);     % initialize iterative solution
  
iterations = 20;  % do a fixed number of iterations

% todo - iterate until specified error tolerance reached or break
%        if we exceed a specified maximum number of iterations

for i=1:iterations
  x_next = ld_inv_b - ld_inv_u*x;
  x = x_next;
end








